#! /usr/bin/env python
#coding=utf-8

import os
import sys

sys.path.append(os.path.join(os.path.dirname(os.path.realpath(__file__)), "../"))
from macho import MachOFileDeps
from tablebuilder import SymbolsTableBuilder, DepedenceiesTableBuilder
from tablebuilder import UndefinedSymbolsTableBuilder, ModuleTableBuilder

def symbol_preprocessor(symbol, idx, module):
    symbol["parent_id"] = module["id"]
    symbol["id"] = symbol.get_id(idx)

def deps_preprocessor(dep, idx, cursor):
    dep.updateCallsFromDb(cursor)

class DylibFile(MachOFileDeps):
    def __init__(self, file, prefix, mgr):
        super(DylibFile, self).__init__(file, prefix, mgr)

    def createSymbolTables(_cursor):
        builder = SymbolsTableBuilder(_cursor)
        builder.createTable()

        builder = UndefinedSymbolsTableBuilder(_cursor)
        builder.createTable()

    def queryCalls(self, _cursor):
        print("Query calls for %d:%s now ..." % (self["id"], self["name"]))
        for dep in self["deps"]:
            dep.queryCalls(_cursor)

    def addCalls(self, _cursor):
        print("Add calls for %d:%s now ..." % (self["id"], self["name"]))
        for dep in self["deps"]:
            dep.addMatchedCallsToDb(_cursor)

    def addDeps(self, _cursor):
        print("Add deps for %d:%s now ..." % (self["id"], self["name"]))
        builder = DepedenceiesTableBuilder(_cursor)
        builder.addMultiObjToDb(self["deps"], _cursor, deps_preprocessor)

    def addSymbols(self, _cursor, logFile=None):
        provided, undefined = self.symbols(logFile)
        print("Add symbols for %d:%s [%d:%d] now ..." % (self["id"], self["name"], len(provided), len(undefined)))
        if logFile:
            logFile.write("Add %d symbols %d undefines for %d:%s now ...\n" % (len(provided), len(undefined), self["id"], self["name"]))
        builder = SymbolsTableBuilder(_cursor)
        builder.addMultiObjToDb(provided, self, symbol_preprocessor, logFile)

        if logFile:
            logFile.write("Add %d undefines for %d:%s now ...\n" % (len(undefined), self["id"], self["name"]))

        builder = UndefinedSymbolsTableBuilder(_cursor)
        builder.addMultiObjToDb(undefined, self, symbol_preprocessor, logFile)

        return (len(provided), len(undefined))

    # All provided symbols
    def __get_provided_symbols_cnt(self, cursor):
        if "provided" in self:
            return self["provided"]
        try:
            sqlcmd = "select count(id) from symbols where parent_id=%d" % (self["id"])
            cursor.execute(sqlcmd)
            for row in cursor:
                return int(row[0])
        except:
            return 0

    def __get_public_symbols_cnt(self, cursor):
        if "public_symbols" in self:
            return self["public_symbols"]
        try:
            sqlcmd = "select count(id) from symbols where parent_id=%d and bits!='l'" % (self["id"])
            cursor.execute(sqlcmd)
            for row in cursor:
                return int(row[0])
        except:
            return 0

    # Provided Symbols being used
    def __get_used_symbols_cnt(self, cursor):
        try:
            sqlcmd = "select count(distinct(symbol_id)) from calls where callee_id=%d" % (self["id"])
            cursor.execute(sqlcmd)
            for row in cursor:
                return int(row[0])
        except:
            return 0

    # Get needed symbols cnt
    def get_needed_symbols_cnt(self, cursor):
        try:
            sqlcmd = "select count(name) from undefines where parent_id=%d" % (self["id"])
            cursor.execute(sqlcmd)
            for row in cursor:
                return int(row[0])
        except:
            return 0

    # Undefined symbols being provided by depended libraries
    def __get_matched_symbols_cnt(self, cursor):
        try:
            sqlcmd = "select count(distinct(symbol_id)) from calls where caller_id=%d" % (self["id"])
            cursor.execute(sqlcmd)
            for row in cursor:
                return int(row[0])
        except:
            return 0

    # Undefined symbols duplicated by depended libraries
    def __get_duplicated_symbols_cnt(self, cursor):
        try:
            sqlcmd = "select count(name) from call_details where caller_id=%d GROUP by name HAVING count(*)  > 1" % (self["id"])
            cursor.execute(sqlcmd)
            for row in cursor:
                return int(row[0])
        except:
            return 0

    # Undefined symbols being provided by depended libraries
    def __get_unmatched_symbols_cnt(self, cursor):
        try:
            sqlcmd = "select count(name) from undefines where parent_id=%d and name not in (select distinct(name) from call_details where caller_id=%d)" % (self["id"], self["id"])
            cursor.execute(sqlcmd)
            for row in cursor:
                return int(row[0])
        except:
            return 0

    def updateCountsInfo(self, cursor):
        self["needed"] = self.get_needed_symbols_cnt(cursor)
        self["matched"] = self.__get_matched_symbols_cnt(cursor)
        self["duplicated"] = 0#self.__get_duplicated_symbols_cnt(cursor)
        self["unmatched"] = self.__get_unmatched_symbols_cnt(cursor)

        self["provided"] = self.__get_provided_symbols_cnt(cursor)
        self["used"] = self.__get_used_symbols_cnt(cursor)
        self["public_symbols"] = self.__get_public_symbols_cnt(cursor)

if __name__ == '__main__':
    import sqlite3

    _conn = sqlite3.connect(os.path.join(os.path.dirname(os.path.realpath(__file__)), "archinfo.db"))
    _cursor = _conn.cursor()

    from macho import MachOFileMgr
    from dependency import Dependency
    mgr = MachOFileMgr(fileClass = DylibFile, dependenceClass=Dependency)
    mgr.scan_all_files()
    #dylib = mgr.get_file_by_path("/usr/lib/libutil.dylib")
    dylib = mgr.get_file_by_path("/System/Library/PrivateFrameworks/SAObjects.framework/Versions/A/SAObjects")

    DylibFile.createSymbolTables(_cursor)
    with open(os.path.join(os.path.dirname(os.path.realpath(__file__)), "add_symbols.log"), "w") as f:
        dylib.addSymbols(_cursor, f)
    #dylib.addCalls(_cursor)
    #dylib.addDeps(_cursor)

    _conn.commit()
    _cursor.close()
    _conn.close()
